#!/usr/bin/env python3
import rospy
import threading
import numpy as np
import torch
from func.f_ego_motion_comp import EgoMotionComp
from uav_event.msg import CameraData, UInt32Mat
from rich import print


publisherInitialized = False
nodePub = None


subscriberInitialized = False
nodeSub = None


scale = 4
snn_scale = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
femc = EgoMotionComp(scale = scale, device = device)


sumTime = 0.0
counter = -20
minDuration = torch.inf
maxDuration = 0


def nodeSubCallback(msg: CameraData) -> None:
    global femc, publisherInitialized, nodePub, device, scale, snn_scale, sumTime, counter, minDuration, maxDuration
    msgTime = msg.header.stamp.secs + msg.header.stamp.nsecs * 1e-9

    imuData = torch.tensor(msg.imus, dtype = torch.float, device = device)
    if imuData.shape[0]:
        imuSize = len(msg.imucol)
        imuData = imuData.reshape(imuData.shape[0] // imuSize, imuSize) # [t, wz, wy, wx]
    femc.updateIMUTime(msgTime)
    femc.processImu(imuData[-1:])

    eventData = torch.tensor(msg.events, dtype = torch.long, device = device)
    if eventData.shape[0]:
        eventSize = len(msg.eventcol)
        eventData = eventData.reshape(eventData.shape[0] // eventSize, eventSize) # [sec, nsec, p, y, x]
        eventData[:, 3:] = torch.floor(eventData[:, 3:] / scale)

    res = femc.execute(eventData)

    if res.shape[0]:
        res[:, 3:] = torch.floor(res[:, 3:] / (snn_scale / scale))

    pubRes = UInt32Mat()
    pubRes.header.stamp.secs = msg.header.stamp.secs
    pubRes.header.stamp.nsecs = msg.header.stamp.nsecs
    pubRes.data = res.flatten().cpu().numpy().tolist()
    pubRes.column = msg.eventcol
    if publisherInitialized:
        nodePub.publish(pubRes)

    if counter > 0:
        duration = sum(femc.processTimes)
        sumTime += duration
        if minDuration > duration:
            minDuration = duration
        if maxDuration < duration:
            maxDuration = duration
        print("Ego: min %.6f ms, max %.6f ms, avr %.6f ms" % (1000 * minDuration, 1000 * maxDuration, 1000 * sumTime / counter))
    counter += 1

def subThreadFunc() -> None:
    global subscriberInitialized, nodeSub
    nodeSub = rospy.Subscriber("/uav_event/event_imu_data", CameraData, nodeSubCallback)
    subscriberInitialized = True
    rospy.spin()


def pubThreadFunc() -> None:
    global publisherInitialized, nodePub
    nodePub = rospy.Publisher("/uav_event/ego_motion_comp_res", UInt32Mat, queue_size = 1)
    publisherInitialized = True


if __name__ == "__main__":
    rospy.init_node("ego_motion_comp_controller", anonymous = True)

    subT = threading.Thread(target = subThreadFunc)
    pubT = threading.Thread(target = pubThreadFunc)

    subT.start()
    pubT.start()

    subT.join()
    pubT.join()
